import argparse
import sys

# argv = sys.argv
# 数据集名称从运行命令中获得
# dataset = argv[1]
# print(f"Dataset: {dataset}")
# 这里先不设置默认为ACM
# dataset = "ACM"

# 在这个脚本中，parse_arguments() 函数用来解析命令行参数，并从中获取数据集的名称。
# set_params(dataset) 函数根据传入的数据集名称调用相应的参数设置函数。
# 最终，__main__ 中的逻辑会根据命令行输入设置相应的数据集参数并打印出来。

def parse_args(args=None):
    parser = argparse.ArgumentParser(description='SeHGNN')

    ## For environment costruction


    return parser.parse_args(args)


def dblp_params(args=None):
    parser = argparse.ArgumentParser(description='HINormer')
    parser.add_argument('--feats-type', type=int, default=3,
                    help='Type of the node features used. ' +
                         '0 - loaded features; ' +
                         '1 - only target node features (zero vec for others); ' +
                         '2 - only target node features (id vec for others); ' +
                         '3 - all id vec. Default is 2' +
                         '4 - only term features (id vec for others);' +
                         '5 - only term features (zero vec for others).')
    parser.add_argument('--device', type=int, default=0)
    parser.add_argument('--hidden-dim', type=int, default=256,
                    help='Dimension of the node hidden state. Default is 32.')
    parser.add_argument('--dataset', type=str, default='DBLP', help='DBLP, IMDB, Freebase, AMiner, DBLP-HGB, IMDB-HGB')
    parser.add_argument('--num-heads', type=int, default=2,
                    help='Number of the attention heads. Default is 2.')
    parser.add_argument('--epoch', type=int, default=1000, help='Number of epochs.')
    parser.add_argument('--patience', type=int, default=50, help='Patience.')
    parser.add_argument('--repeat', type=int, default=5, help='Repeat the training and testing for N times. Default is 1.')
    parser.add_argument('--num-layers', type=int, default=2, help='The number of layers of HINormer layer')
    parser.add_argument('--num-gnns', type=int, default=4,
                    help='The number of layers of both structural and heterogeneous encoder')
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--seed', type=int, default=2023)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--weight-decay', type=float, default=0)
    parser.add_argument('--len-seq', type=int, default=50, help='The length of node sequence.')
    parser.add_argument('--l2norm', type=bool, default=True, help='Use l2 norm for prediction')
    parser.add_argument('--mode', type=int, default=0,
                    help='Output mode, 0 for offline evaluation and 1 for online HGB evaluation')
    parser.add_argument('--temperature', type=float, default=2.0, help='Temperature of attention score')
    parser.add_argument('--beta', type=float, default=0.1, help='Weight of heterogeneity-level attention score')

    return parser.parse_args(args)


def init_params(parser):
    # For environment costruction
    parser.add_argument('--device', type=int, default=0)